//////////////////////////////////////////////////////////////////////
// SparseMatrix.h
//
// SHEN Fangyang
// me@shenfy.com
//
// Copyright (C) SHEN Fangyang, 2011, All rights reserved.
//////////////////////////////////////////////////////////////////////

//////////////////////////////////////////////////////////////////////
//  Brief Description
//
//	SIMPLE LINEAR ALGEBRA PACKAGE
//
//	Sparse matrix in MKL CSR/CSC format, Zero-based, as in
//	Intel(R) MKL 10.0 Manual Appendix A, page 2875
//
//////////////////////////////////////////////////////////////////////
#pragma once

#include "FileMultiSec.h"
#include "SysTypeTraits.h"
#include <ostream>
#include "meta.h"
#include "LA_basic.h"
#include <mkl_spblas.h>

namespace grflib
{
	namespace LA
	{
		//////////////////////////////////////////////////////////
		// Declaration
		//
		// 2ndd template parameter CSR:
		//		true  -->  CSR format
		//		false -->  CSC format
		//////////////////////////////////////////////////////////
		template<typename matrix_type, bool CSR = true>
		struct CSparseMatrix
		{
		public:
			struct MatMeta
			{
				UINT nRow, nCol;
				UINT nNonZero;
				float scale;
			};

			//ctor & dtor
			CSparseMatrix(void);
			virtual ~CSparseMatrix(void);

			//compression & decompression
			bool Compress(const matrix_type& denseMat, typename matrix_type::value_type threshold, UINT nNonZero);
			bool Decompress(matrix_type& denseMat);

			//compressed data
			typename matrix_type::value_type *m_values;
			UINT *m_indices;
			UINT *m_ptrB;
			UINT *m_ptrE;
			MatMeta m_meta;

			//disk IO
			bool WriteToFile(w32lib::CFileRaw *pFile);
			bool LoadFromFile(w32lib::CFileRaw *pFile);
			
			//text dump
			template<typename TT, bool CCSSRR>
			friend std::ostream & operator<< (std::ostream & out, 
											  const CSparseMatrix<TT, CCSSRR>& mat);

			//traits
			static const bool m_isCSR = CSR;
			typedef typename matrix_type::value_type value_type;

		protected:
			bool Alloc(void);
			void Release(void);
		};

		namespace traits
		{
			template<typename MatrixType>
			inline UINT Major(const MatrixType &m) {return (m.m_isCSR ? m.m_meta.nRow : m.m_meta.nCol);}

			template<typename MatrixType>
			inline UINT Minor(const MatrixType &m) {return (m.m_isCSR ? m.m_meta.nCol : m.m_meta.nRow);}
		}

		//////////////////////////////////////////////////////////
		//Definition
		template<typename matrix_type, bool CSR>
		CSparseMatrix<matrix_type, CSR>::CSparseMatrix() :
			m_values(0),
			m_indices(0),
			m_ptrB(0),
			m_ptrE(0)
		{
			memset(&m_meta, 0, sizeof(MatMeta));
		}

		template<typename matrix_type, bool CSR>
		CSparseMatrix<matrix_type, CSR>::~CSparseMatrix()
		{
			Release();
		}

		template<typename matrix_type, bool CSR>
		void CSparseMatrix<matrix_type, CSR>::Release()
		{
			_SafeDelete(m_values);
			_SafeDelete(m_indices);
			_SafeDelete(m_ptrB);
			_SafeDelete(m_ptrE);
			m_values = 0;
			m_indices = 0;
			m_ptrB = 0;
			m_ptrE = 0;
			memset(&m_meta, 0, sizeof(MatMeta));
		}

		template<typename matrix_type, bool CSR>
		bool CSparseMatrix<matrix_type, CSR>::Compress(const matrix_type &denseMat, typename matrix_type::value_type threshold, UINT nNonZero)
		{
			//allocate memory
			if (m_meta.nRow != denseMat.size1() || m_meta.nCol != denseMat.size2())
			{
				Release();
				m_meta.nRow = denseMat.size1(); m_meta.nCol = denseMat.size2();
				m_meta.nNonZero = nNonZero;
				if (!Alloc())
				{
					Release();
					return false;
				}
			}

			//compression
			typename matrix_type::value_type *pDst;
			typename matrix_type::value_type value;
			UINT *pIdx;
			UINT count = 0;
			UINT major = traits::Major(*this);
			UINT minor = traits::Minor(*this);

			for (UINT i1 = 0; i1 < major; i1++)
			{
				*(m_ptrB + i1) = count;
				for (UINT i2 = 0; i2 < minor; i2++)
				{
					__staticIf(CSR) //row major
					{	value = denseMat(i1, i2);	}
					__staticIfNot(CSR) //column major
					{	value = denseMat(i2, i1);	}

					if (abs(value) > threshold) //add to sparse matrix
					{
						pDst = m_values + count;
						*pDst = value;
						pIdx = m_indices + count;
						*pIdx = i2;
						count ++;
					}
				}
				//column ends
				*(m_ptrE + i1) = count;
			}

			return true;
		}

		template<typename matrix_type, bool CSR>
		bool CSparseMatrix<matrix_type, CSR>::Decompress(matrix_type &denseMat)
		{
			if (denseMat.size1() != m_meta.nRow || denseMat.size2() != m_meta.nCol)
				denseMat.resize(m_meta.nRow, m_meta.nCol, false);

			UINT line = 0, count = 0;
			UINT lineBegin, lineNed;
			UINT major = traits::Major(*this);
			
			for (UINT line = 0; line < major; line++)
			{
				lineBegin = m_ptrB[line];
				lineEnd = m_ptrE[line];
				for (UINT idx = lineBegin; idx < lineEnd; idx++)
				{
					__staticIf(CSR)
					{	denseMat(line, m_indices[idx]) = m_values[idx];	}
					__staticIfNot(CSR)
					{	denseMat(m_indices[idx], line) = m_values[idx];	}
				}
			}

			return true;
		}



		template<typename matrix_type, bool CSR>
		bool CSparseMatrix<matrix_type, CSR>::Alloc()
		{
			UINT major = traits::Major(*this);

			m_values = new typename matrix_type::value_type[m_meta.nNonZero];
			m_indices = new UINT[m_meta.nNonZero];
			m_ptrB = new UINT[major];
			m_ptrE = new UINT[major];
			if (!m_values || !m_indices || !m_ptrB || !m_ptrE)
			{
				Release();
				W32ERR(_T("Failed to allocate memory for sparse matrix!"));
				return false;
			}
			return true;
		}

		template<typename matrix_type, bool CSR>
		std::ostream & operator<< (std::ostream &out, const CSparseMatrix<matrix_type, CSR>& mat)
		{
			out << "values = (";
			for (UINT i = 0; i < mat.m_meta.nNonZero; i++)
			{
				out << mat.m_values[i] << " ";
			}
			out << ")" << std::endl;

			out << "rows = (";
			for (UINT i = 0; i < mat.m_meta.nNonZero; i++)
			{
				out << mat.m_indices[i] << " ";
			}
			out << ")" << std::endl;

			UINT major = traits::Major(mat);
			out << "pointerB = (";
			for (UINT i = 0; i < major; i++)
			{
				out << mat.m_ptrB[i] << " ";
			}
			out << ")" << std::endl;

			out << "pointerE = (";
			for (UINT i = 0; i < major; i++)
			{
				out << mat.m_ptrE[i] << " ";
			}
			out << ")" << std::endl;

			return out;
		}

		template<typename matrix_type, bool CSR>
		bool CSparseMatrix<matrix_type, CSR>::WriteToFile(w32lib::CFileRaw *pFile)
		{
			if (!pFile || !pFile->IsOpen())
				return false;

			UINT major = traits::Major(*this);
			pFile->Write(&m_meta, sizeof(MatMeta));
			pFile->Write(m_values, sizeof(typename matrix_type::value_type) * m_meta.nNonZero);
			pFile->Write(m_indices, sizeof(UINT) * m_meta.nNonZero);
			pFile->Write(m_ptrB, sizeof(UINT) * major);
			pFile->Write(m_ptrE, sizeof(UINT) * major);
			return true;
		}

		template<typename matrix_type, bool CSR>
		bool CSparseMatrix<matrix_type, CSR>::LoadFromFile(w32lib::CFileRaw *pFile)
		{
			if (!pFile || !pFile->IsOpen())
				return false;

			pFile->Read(&m_meta, sizeof(MatMeta));
			if (!Alloc())
			{
				Release();
				return false;
			}

			UINT major = traits::Major(*this);
			pFile->Read(m_values, sizeof(typename matrix_type::value_type) * m_meta.nNonZero);
			pFile->Read(m_indices, sizeof(typename matrix_type::value_type) * m_meta.nNonZero);
			pFile->Read(m_ptrB, sizeof(UINT) * major);
			pFile->Read(m_ptrE, sizeof(UINT) * major);

			return true;
		}

	} //namespace LA
} //namespace grflib